import sys
sys.path.append("../")

from new_arm_control.angle import ArmAngle
from new_arm_control.move import ArmMove
from new_arm_control.func import pick_and_release

import multiprocessing as mp
import time
import cv2
import os
from helper import *
import apriltag


CUBE_HEIGHT = 3.6
MAX_NUMBER_OF_OBEJECTS = 6

PICK_AND_RELEASE_OFFSET = 6
ARM_Z_OFFSET = 5

WINDOW_W = 640
WINDOW_H = 480
WINDOW_NAME = "Project I: Apriltag Sort"

CONFIG_FILE = os.path.expanduser("~/.bc_config/cam_calibration.config")
cfg = load_config(CONFIG_FILE)

arm_base_coordinate_x = cfg["arm_base_coordinate"]["x"]
arm_base_coordinate_y = cfg["arm_base_coordinate"]["y"]
pixel_per_cm_ratio = cfg["pixel_per_cm_ratio"]

# 放置区域范围
'''
x轴坐标：机械臂基座在摄像头视野中的位置减去10.5cm乘以每个cm对应的像素数
y轴坐标：机械臂基座在摄像头视野中的位置连续增加6.5cm乘以每个cm对应的像素数
'''
place_position_in_pixels = {
    1: (arm_base_coordinate_x - int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(13.5 * pixel_per_cm_ratio)),
    2: (arm_base_coordinate_x - int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(20 * pixel_per_cm_ratio)),
    3: (arm_base_coordinate_x - int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(26.5 * pixel_per_cm_ratio)),
    4: (arm_base_coordinate_x + int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(13.5 * pixel_per_cm_ratio)),
    5: (arm_base_coordinate_x + int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(20 * pixel_per_cm_ratio)),
    6: (arm_base_coordinate_x + int(10.5 * pixel_per_cm_ratio), arm_base_coordinate_y + int(26.5 * pixel_per_cm_ratio)),
}

def nothing(_):
    pass

# 计算机械臂在摄像头像素坐标下距离原点的实际距离(x,y,z)
arm_base_pos = [
    *compute_pixel_to_dist(
        arm_base_coordinate_x, arm_base_coordinate_y, pixel_per_cm_ratio
    ),
    0,
]

def process_wrapper(queue, signal):
    def decorator(func):
        def wrapper():
            arm_move = ArmMove()
            arm_angle = ArmAngle(base_pos=arm_base_pos)
            while True:
                if not queue.empty():
                    targets = queue.get()
                    func(arm_move, arm_angle, *targets)
                    signal.set()
                time.sleep(0.1)  # Prevents busy waiting
        return wrapper
    return decorator

# Create a multiprocessing queue
queue = mp.Queue()
signal = mp.Event()
signal.set()


# Create a decorated version of pick_and_release
@process_wrapper(queue, signal)
def wrapped_pick_and_release(*args):
    pick_and_release(*args)


def run(min_dist_threshold = 10):
    # Start the process
    process = mp.Process(target=wrapped_pick_and_release)
    process.start()
    time.sleep(3)
    cap = cv2.VideoCapture(0)
    cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(WINDOW_NAME, WINDOW_W, WINDOW_H)

    print("按'q'关闭摄像头")

    place_position_in_dist = {
        name: (*compute_pixel_to_dist(x, y, pixel_per_cm_ratio), CUBE_HEIGHT)
        for name, (x, y) in place_position_in_pixels.items()
    }
    
    print(place_position_in_dist)
    
    target_box_size_in_pixel = int(CUBE_HEIGHT * pixel_per_cm_ratio)

    i = 0
    
    detector = apriltag.Detector()
    
    try:
        ### 整个while 是检测函数的主体, 在主体里面插入 pick_and_release 函数进程的控制信号量
        while True:
            ret, frame = cap.read()
            gray_image = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            # 对输入的图像检测 Apriltag
            tag_pos_pairs = []
            results = detector.detect(gray_image)
            
            for result in results:
                tag_id = result.tag_id
                top_left, top_right, bottom_right, bottom_left = result.corners
                center_x = int((top_left[0] + bottom_right[0]) / 2)
                center_y = int((top_left[1] + bottom_right[1]) / 2)
                tag_inform = (tag_id, center_x, center_y)
                tag_pos_pairs.append(tag_inform)

            # 可视化检测结果
            for tag in tag_pos_pairs:
                draw_tag(frame, tag[0], tag[1:])

            # 可视化放置位置
            frame = draw_target_pos(frame, place_position_in_pixels, width=target_box_size_in_pixel)
            
            cv2.putText(
                frame,
                f"Jobs Count: {i}",
                (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX,
                1,
                (0, 255, 0),
                2,
            )

            # 当前正在抓取中,不执行任何其它抓取判断
            if not signal.is_set():
                time.sleep(0.03)

            else:
                ##  控制信号量
                tag_pos_pairs_sorted_by_id = sorted(tag_pos_pairs, key=lambda x: x[0])

                # 判断检测到的tags 是否都在对应的位置, 返回第一个不对应的tag
                mismatched_id, x, y = mismatched_position(
                    place_position_in_pixels, tag_pos_pairs_sorted_by_id, min_dist_threshold
                )

                if mismatched_id != -1:
                    mismatched_apriltag_coordinate = compute_pixel_to_dist(
                        x, y, pixel_per_cm_ratio
                    )
                    
                    print(f"将坐标发送到队列,  当前放置的Apriltag为 {mismatched_id}  {mismatched_apriltag_coordinate[0]}  {mismatched_apriltag_coordinate[1]}")
                    
                    queue.put(
                        (
                            mismatched_apriltag_coordinate[0],
                            mismatched_apriltag_coordinate[1],
                            CUBE_HEIGHT,
                            place_position_in_dist[mismatched_id][0],
                            place_position_in_dist[mismatched_id][1],
                            place_position_in_dist[mismatched_id][2],
                            PICK_AND_RELEASE_OFFSET,
                            PICK_AND_RELEASE_OFFSET
                        )
                    )

                    i += 1
                    signal.clear()
            cv2.imshow("Video Frame", frame)
            if cv2.waitKey(10) & 0xFF == ord('q'):
                break

    except KeyboardInterrupt:
        print("Process interrupted by user")
    finally:
        # Ensure the process is properly terminated
        process.terminate()
        process.join()


if __name__ == "__main__":
    
    min_dist_threshold = 15
    
    run(min_dist_threshold)
